#include "main/utils.h"
#include "lib/binary_tree.h"
using namespace std;

using TreeNode = BinaryTreeNode;

class Solution {
  private:
    TreeNode* ans;
    unordered_map<int, TreeNode*> fa;
    unordered_map<int, bool> vis;
  public:
    bool dfs(TreeNode* root, TreeNode* p, TreeNode* q) {
      if (!root) {
        return false;
      }
      bool lson = dfs(root->left, p, q);
      bool rson = dfs(root->right, p, q);
      if ((lson && rson) || ((root->value == p->value || root->value == q->value) && (lson || rson))) {
        ans = root;
      }
      return lson || rson || (root->value == p->value) || (root->value == q->value);
    }
    TreeNode* lowestCommonAncestor(TreeNode* root, TreeNode* p, TreeNode* q) {
      dfs(root, p, q);
      return ans;
    }

    void dfs(TreeNode* root) {
      if (root->left) {
        fa[root->left->value] = root;
        dfs(root->left);
      }
      if (root->right) {
        fa[root->right->value] = root;
        dfs(root->right);
      }
    }
    TreeNode* lowestCommonAncestor1(TreeNode* root, TreeNode* p, TreeNode* q) {
      fa[root->value] = nullptr;
      dfs(root);
      while (p) {
        vis[p->value] = 1;
        p = fa[p->value];
      }
      while (q) {
        if (vis[q->value]) {
          return q;
        }
        q = fa[q->value];
      }
      return nullptr;
    }
};

int main() {
  BinaryTreeNode* node1 = CreateBinaryTreeNode(1);
  BinaryTreeNode* node2 = CreateBinaryTreeNode(2);
  BinaryTreeNode* node3 = CreateBinaryTreeNode(3);
  BinaryTreeNode* node4 = CreateBinaryTreeNode(4);
  BinaryTreeNode* node5 = CreateBinaryTreeNode(5);
  BinaryTreeNode* node6 = CreateBinaryTreeNode(6);
  BinaryTreeNode* node7 = CreateBinaryTreeNode(7);
  BinaryTreeNode* node8 = CreateBinaryTreeNode(8);
  BinaryTreeNode* node9 = CreateBinaryTreeNode(9);

  ConnectNode(node1, node2, node3);
  ConnectNode(node2, node4, node5);
  ConnectNode(node4, node6, node7);
  ConnectNode(node5, node8, node9);

  Solution su;
  auto ret = su.lowestCommonAncestor(node1, node6, node8);
  cout << "The lowest common binary node: " << ret->value << endl;
  auto ret1 = su.lowestCommonAncestor1(node1, node6, node8);
  cout << "The lowest common binary node: " << ret1->value << endl;

  return 0;
}
